import torch.nn as nn
import torch


#------------------------#
#           G
#------------------------#
class Generator(nn.Module):
    def __init__(self, z_dim):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(z_dim + 10, 128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(in_features=512, out_features=28 * 28),
            nn.Tanh()
        )

    def forward(self, z, c):
        # print(z.shape,c.shape)
        x = self.model(torch.cat([z, c], dim=1))
        x = x.view(-1, 1, 28, 28)
        return x



#--------------------------------#
#               D
#--------------------------------#
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(28 * 28 + 10, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, x, c):
        x = x.view(x.size(0), -1)
        validity = self.model(torch.cat([x, c], -1))
        return validity